"""
声调曲线拟合模块
根据声调类型进行一次或二次函数拟合
"""
import numpy as np
from scipy.optimize import curve_fit
from typing import Tuple, List, Dict, Optional
import warnings

class ToneCurveFitter:
    def __init__(self):
        self.fitting_methods = {
            1: self._fit_linear,    # 一声：一次函数
            2: self._fit_linear,    # 二声：一次函数  
            3: self._fit_quadratic, # 三声：二次函数
            4: self._fit_linear     # 四声：一次函数
        }
    
    def _linear_function(self, x: np.ndarray, a: float, b: float) -> np.ndarray:
        """一次函数: y = ax + b"""
        return a * x + b
    
    def _quadratic_function(self, x: np.ndarray, a: float, b: float, c: float) -> np.ndarray:
        """二次函数: y = ax² + bx + c"""
        return a * x**2 + b * x + c
    
    def _fit_linear(self, time_array: np.ndarray, f0_array: np.ndarray) -> Dict:
        """
        一次函数拟合
        
        Args:
            time_array: 时间序列
            f0_array: 基频序列
            
        Returns:
            拟合结果字典
        """
        try:
            # 标准化时间到[0,1]区间
            time_normalized = (time_array - time_array.min()) / (time_array.max() - time_array.min())
            
            # 拟合一次函数
            popt, pcov = curve_fit(self._linear_function, time_normalized, f0_array)
            a, b = popt
            
            # 计算拟合优度
            y_pred = self._linear_function(time_normalized, a, b)
            r_squared = self._calculate_r_squared(f0_array, y_pred)
            
            # 计算拟合误差
            rmse = np.sqrt(np.mean((f0_array - y_pred)**2))
            
            return {
                'type': 'linear',
                'parameters': {'a': a, 'b': b},
                'function': self._linear_function,
                'r_squared': r_squared,
                'rmse': rmse,
                'success': True,
                'time_normalized': time_normalized,
                'fitted_values': y_pred
            }
        except Exception as e:
            print(f"一次函数拟合失败: {e}")
            return {
                'type': 'linear',
                'parameters': {'a': 0, 'b': np.mean(f0_array) if len(f0_array) > 0 else 0},
                'function': self._linear_function,
                'r_squared': 0,
                'rmse': float('inf'),
                'success': False,
                'error': str(e)
            }
    
    def _fit_quadratic(self, time_array: np.ndarray, f0_array: np.ndarray) -> Dict:
        """
        二次函数拟合
        
        Args:
            time_array: 时间序列
            f0_array: 基频序列
            
        Returns:
            拟合结果字典
        """
        try:
            # 标准化时间到[0,1]区间
            time_normalized = (time_array - time_array.min()) / (time_array.max() - time_array.min())
            
            # 拟合二次函数
            popt, pcov = curve_fit(self._quadratic_function, time_normalized, f0_array)
            a, b, c = popt
            
            # 计算拟合优度
            y_pred = self._quadratic_function(time_normalized, a, b, c)
            r_squared = self._calculate_r_squared(f0_array, y_pred)
            
            # 计算拟合误差
            rmse = np.sqrt(np.mean((f0_array - y_pred)**2))
            
            return {
                'type': 'quadratic',
                'parameters': {'a': a, 'b': b, 'c': c},
                'function': self._quadratic_function,
                'r_squared': r_squared,
                'rmse': rmse,
                'success': True,
                'time_normalized': time_normalized,
                'fitted_values': y_pred
            }
        except Exception as e:
            print(f"二次函数拟合失败: {e}")
            return {
                'type': 'quadratic',
                'parameters': {'a': 0, 'b': 0, 'c': np.mean(f0_array) if len(f0_array) > 0 else 0},
                'function': self._quadratic_function,
                'r_squared': 0,
                'rmse': float('inf'),
                'success': False,
                'error': str(e)
            }
    
    def _calculate_r_squared(self, y_true: np.ndarray, y_pred: np.ndarray) -> float:
        """计算R²决定系数"""
        ss_res = np.sum((y_true - y_pred) ** 2)
        ss_tot = np.sum((y_true - np.mean(y_true)) ** 2)
        
        if ss_tot == 0:
            return 1.0 if ss_res == 0 else 0.0
        
        return 1 - (ss_res / ss_tot)
    
    def fit_tone_curve(self, time_array: np.ndarray, f0_array: np.ndarray, tone: int) -> Dict:
        """
        根据声调类型拟合曲线
        
        Args:
            time_array: 时间序列
            f0_array: 基频序列
            tone: 声调 (1-4)
            
        Returns:
            拟合结果字典
        """
        if len(time_array) == 0 or len(f0_array) == 0:
            return {
                'type': 'unknown',
                'parameters': {},
                'success': False,
                'error': '输入数据为空'
            }
        
        if len(time_array) != len(f0_array):
            return {
                'type': 'unknown', 
                'parameters': {},
                'success': False,
                'error': '时间序列和基频序列长度不匹配'
            }
        
        if tone not in self.fitting_methods:
            return {
                'type': 'unknown',
                'parameters': {},
                'success': False,
                'error': f'不支持的声调: {tone}'
            }
        
        # 调用对应的拟合方法
        fitting_method = self.fitting_methods[tone]
        result = fitting_method(time_array, f0_array)
        result['tone'] = tone
        
        return result
    
    def generate_sample_points(self, fitting_result: Dict, num_points: int = 20) -> np.ndarray:
        """
        从拟合曲线均匀采样指定数量的点
        
        Args:
            fitting_result: 拟合结果字典
            num_points: 采样点数量
            
        Returns:
            采样点的基频值数组
        """
        if not fitting_result.get('success', False):
            return np.array([])
        
        # 在[0,1]区间均匀采样
        x_sample = np.linspace(0, 1, num_points)
        
        # 根据拟合类型计算对应的y值
        if fitting_result['type'] == 'linear':
            params = fitting_result['parameters']
            y_sample = self._linear_function(x_sample, params['a'], params['b'])
        elif fitting_result['type'] == 'quadratic':
            params = fitting_result['parameters']
            y_sample = self._quadratic_function(x_sample, params['a'], params['b'], params['c'])
        else:
            return np.array([])
        
        return y_sample
    
    def evaluate_fitting_quality(self, fitting_result: Dict) -> str:
        """
        评估拟合质量
        
        Args:
            fitting_result: 拟合结果字典
            
        Returns:
            拟合质量描述
        """
        if not fitting_result.get('success', False):
            return "拟合失败"
        
        r_squared = fitting_result.get('r_squared', 0)
        
        if r_squared >= 0.9:
            return "优秀"
        elif r_squared >= 0.8:
            return "良好"
        elif r_squared >= 0.6:
            return "一般"
        else:
            return "较差"

# 使用示例
if __name__ == "__main__":
    fitter = ToneCurveFitter()
    
    # 模拟测试数据
    time_test = np.linspace(0, 1, 50)
    
    # 测试一声（上升调）
    f0_tone1 = 200 + 50 * time_test + np.random.normal(0, 5, 50)
    result1 = fitter.fit_tone_curve(time_test, f0_tone1, 1)
    print("一声拟合结果:")
    print(f"  类型: {result1['type']}")
    print(f"  参数: {result1['parameters']}")
    print(f"  R²: {result1.get('r_squared', 0):.3f}")
    print(f"  质量: {fitter.evaluate_fitting_quality(result1)}")
    
    # 生成采样点
    samples1 = fitter.generate_sample_points(result1)
    print(f"  采样点数: {len(samples1)}")
    
    print()
    
    # 测试三声（低降升调）
    f0_tone3 = 200 - 30 * time_test + 80 * time_test**2 + np.random.normal(0, 5, 50)
    result3 = fitter.fit_tone_curve(time_test, f0_tone3, 3)
    print("三声拟合结果:")
    print(f"  类型: {result3['type']}")
    print(f"  参数: {result3['parameters']}")
    print(f"  R²: {result3.get('r_squared', 0):.3f}")
    print(f"  质量: {fitter.evaluate_fitting_quality(result3)}")
    
    # 生成采样点
    samples3 = fitter.generate_sample_points(result3)
    print(f"  采样点数: {len(samples3)}")
